import torch
from torch.utils.data import DataLoader
import torch.nn as nn
from datasets.dataset import TicketDataset
from models.MPDnnModel import PricePredictor
from utils.preprocess import preprocess
from utils.concat_csv_files import concat_csv_files
from tqdm import tqdm
import json


def evaluate(model_path, test_df, batch_size=32):
    print(f"\n🔎 正在评估模型：{model_path}")

    # 加载模型结构参数
    with open("checkpoints/model_config.json", "r") as f:
        config = json.load(f)

    model = PricePredictor(
        num_numeric_features=config["num_numeric_features"],
        num_film_ids=config["num_film_ids"],
        film_emb_dim=config["film_emb_dim"],
        hidden_dims=config["hidden_dims"],
        dropout=config["dropout"]
    )

    model.load_state_dict(torch.load("checkpoints/best_model.pt"))
    model.eval()

    # 构建 Dataset 和 DataLoader（不从 csv）
    test_dataset = TicketDataset(test_df)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    # 损失函数
    criterion = nn.MSELoss()
    mae_loss = nn.L1Loss()

    total_mse = 0.0
    total_mae = 0.0
    preds = []
    trues = []

    with torch.no_grad():
        for x_numeric, x_film_id, y in tqdm(test_loader, desc="🧪 Testing"):
            output = model(x_numeric, x_film_id)
            loss_mse = criterion(output, y)
            loss_mae = mae_loss(output, y)

            total_mse += loss_mse.item() * x_numeric.size(0)
            total_mae += loss_mae.item() * x_numeric.size(0)

            preds.extend(output.squeeze().tolist())
            trues.extend(y.squeeze().tolist())

    mse = total_mse / len(test_loader.dataset)
    mae = total_mae / len(test_loader.dataset)

    print(f"📈 测试集 MSE: {mse:.4f} | MAE: {mae:.4f}")
    return preds, trues


if __name__ == "__main__":
    TRAIN_DIR = "data/train"
    TEST_DIR = "data/test"
    # 重新读取并预处理原始测试集
    _, test_raw = concat_csv_files(TRAIN_DIR), concat_csv_files(TEST_DIR)
    _, test_df = preprocess(_, test_raw, save_dir="data")

    # 调用测试函数
    evaluate("checkpoints/best_model.pt", test_df)
